"""
Boundary Condition Handling

Implements boundary conditions for the lid-driven cavity problem on L-shaped domains,
including no-slip walls, moving lid, and proper treatment of the reentrant corner.
"""

"""
    BoundaryConditionType

Enumeration of different boundary condition types in the L-shaped domain.
"""
@enum BoundaryConditionType begin
    NO_SLIP_WALL      # Stationary no-slip boundary (u = v = 0)
    MOVING_LID        # Moving lid boundary (u = U_lid, v = 0) 
    CUT_BOUNDARY      # Internal cut boundary from L-shape
    REENTRANT_CORNER  # Special treatment at reentrant corner
end

"""
    apply_boundary_conditions!(ψ, ω, u, v, domain::LShapeDomain{2}, 
                              x::Vector{Float64}, y::Vector{Float64}, 
                              U_lid::Float64, h::Float64)

Apply boundary conditions for 2D lid-driven cavity problem.

# Arguments
- `ψ::Matrix{Float64}`: Stream function field (modified in-place)
- `ω::Matrix{Float64}`: Vorticity field (modified in-place)  
- `u::Matrix{Float64}`: x-velocity field (modified in-place)
- `v::Matrix{Float64}`: y-velocity field (modified in-place)
- `domain::LShapeDomain{2}`: Domain geometry
- `x, y::Vector{Float64}`: Grid coordinates
- `U_lid::Float64`: Moving lid velocity
- `h::Float64`: Grid spacing
"""
function apply_boundary_conditions!(ψ::Matrix{Float64}, ω::Matrix{Float64}, 
                                   u::Matrix{Float64}, v::Matrix{Float64},
                                   domain::LShapeDomain{2}, 
                                   x::Vector{Float64}, y::Vector{Float64},
                                   U_lid::Float64, h::Float64)
    nx, ny = size(ψ)
    
    # Apply stream function boundary conditions (ψ = 0 on all boundaries)
    apply_stream_function_bc_2d!(ψ, domain, x, y, U_lid)
    
    # Apply velocity boundary conditions  
    apply_velocity_bc_2d!(u, v, domain, x, y, U_lid)
    
    # Apply vorticity boundary conditions (derived from no-slip condition)
    apply_vorticity_bc_2d!(ω, ψ, domain, x, y, U_lid, h)
    
    # Set solid region to zero/NaN for visualization
    apply_solid_region_conditions!(ψ, ω, u, v, domain)
end

"""
    apply_stream_function_bc_2d!(ψ::Matrix{Float64}, domain::LShapeDomain{2},
                                x::Vector{Float64}, y::Vector{Float64}, U_lid::Float64)

Apply stream function boundary conditions.
For the standard formulation, ψ = 0 on all boundaries.
"""
function apply_stream_function_bc_2d!(ψ::Matrix{Float64}, domain::LShapeDomain{2},
                                     x::Vector{Float64}, y::Vector{Float64}, 
                                     U_lid::Float64)
    nx, ny = size(ψ)
    
    # ψ = 0 on all boundaries (including moving lid)
    @inbounds for j in 1:ny, i in 1:nx
        if domain.boundary_mask[i, j]
            ψ[i, j] = 0.0
        end
    end
end

"""
    apply_velocity_bc_2d!(u::Matrix{Float64}, v::Matrix{Float64}, domain::LShapeDomain{2},
                         x::Vector{Float64}, y::Vector{Float64}, U_lid::Float64)

Apply velocity boundary conditions.
"""
function apply_velocity_bc_2d!(u::Matrix{Float64}, v::Matrix{Float64}, 
                              domain::LShapeDomain{2},
                              x::Vector{Float64}, y::Vector{Float64}, 
                              U_lid::Float64)
    nx, ny = size(u)
    eps = 1e-12
    
    @inbounds for j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j]
            # Solid region
            u[i, j] = 0.0
            v[i, j] = 0.0
        elseif domain.boundary_mask[i, j]
            # Check if on moving lid
            if is_moving_lid_node(x[i], y[j], Val(2))
                u[i, j] = U_lid  
                v[i, j] = 0.0
            else
                # No-slip wall
                u[i, j] = 0.0
                v[i, j] = 0.0
            end
        end
    end
end

"""
    apply_vorticity_bc_2d!(ω::Matrix{Float64}, ψ::Matrix{Float64}, domain::LShapeDomain{2},
                          x::Vector{Float64}, y::Vector{Float64}, 
                          U_lid::Float64, h::Float64)

Apply vorticity boundary conditions using the no-slip condition.
"""
function apply_vorticity_bc_2d!(ω::Matrix{Float64}, ψ::Matrix{Float64}, 
                               domain::LShapeDomain{2},
                               x::Vector{Float64}, y::Vector{Float64},
                               U_lid::Float64, h::Float64)
    nx, ny = size(ω)
    
    @inbounds for j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j] || !domain.boundary_mask[i, j]
            continue  # Skip solid nodes and interior nodes
        end
        
        # Compute vorticity from no-slip condition: ω = -∂²ψ/∂n²
        ω[i, j] = compute_boundary_vorticity_2d(ψ, i, j, x, y, U_lid, h, domain)
    end
end

"""
    compute_boundary_vorticity_2d(ψ::Matrix{Float64}, i::Int, j::Int,
                                  x::Vector{Float64}, y::Vector{Float64}, 
                                  U_lid::Float64, h::Float64, 
                                  domain::LShapeDomain{2}) -> Float64

Compute vorticity at a boundary node using finite differences.
"""
function compute_boundary_vorticity_2d(ψ::Matrix{Float64}, i::Int, j::Int,
                                      x::Vector{Float64}, y::Vector{Float64},
                                      U_lid::Float64, h::Float64, 
                                      domain::LShapeDomain{2})
    nx, ny = size(ψ)
    eps = 1e-12
    
    # Check which type of boundary this is
    if is_moving_lid_node(x[i], y[j], Val(2))
        # Moving lid: ω = 2(ψ_interior - ψ_wall)/(h²) + 2*U_lid/h
        # Since ψ_wall = 0, and ψ should be computed from interior
        if j > 1 && domain.fluid_mask[i, j-1]
            # Node below is fluid
            return -2.0 * ψ[i, j-1] / (h^2) - 2.0 * U_lid / h
        else
            return -2.0 * U_lid / h  # Fallback
        end
    else
        # No-slip wall: ω = -2*ψ_interior/h²  
        ω_boundary = 0.0
        
        # Check neighbors and use available interior values
        if i > 1 && domain.fluid_mask[i-1, j] && !domain.boundary_mask[i-1, j]
            ω_boundary += -2.0 * ψ[i-1, j] / (h^2)
        end
        if i < nx && domain.fluid_mask[i+1, j] && !domain.boundary_mask[i+1, j]  
            ω_boundary += -2.0 * ψ[i+1, j] / (h^2)
        end
        if j > 1 && domain.fluid_mask[i, j-1] && !domain.boundary_mask[i, j-1]
            ω_boundary += -2.0 * ψ[i, j-1] / (h^2)
        end
        if j < ny && domain.fluid_mask[i, j+1] && !domain.boundary_mask[i, j+1]
            ω_boundary += -2.0 * ψ[i, j+1] / (h^2)
        end
        
        return ω_boundary
    end
end

"""
    apply_solid_region_conditions!(ψ, ω, u, v, domain::LShapeDomain{2})

Set fields to zero (or NaN) in solid regions for proper visualization.
"""
function apply_solid_region_conditions!(ψ::Matrix{Float64}, ω::Matrix{Float64}, 
                                       u::Matrix{Float64}, v::Matrix{Float64},
                                       domain::LShapeDomain{2})
    nx, ny = size(ψ)
    
    @inbounds for j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j]
            # Set solid region to NaN for visualization  
            ψ[i, j] = NaN
            ω[i, j] = NaN
            u[i, j] = 0.0  # Keep velocities as zero rather than NaN
            v[i, j] = 0.0
        end
    end
end

# 3D versions of boundary condition functions

"""
    apply_boundary_conditions!(ψ, ω, u, v, w, domain::LShapeDomain{3},
                              x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                              U_lid::Float64, h::Float64)

Apply boundary conditions for 3D lid-driven cavity problem.
"""
function apply_boundary_conditions!(ψ::Array{Float64,4}, ω::Array{Float64,4},
                                   u::Array{Float64,3}, v::Array{Float64,3}, 
                                   w::Array{Float64,3},
                                   domain::LShapeDomain{3},
                                   x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                                   U_lid::Float64, h::Float64)
    # Apply stream function boundary conditions for vector components
    apply_vector_stream_function_bc_3d!(ψ, domain, x, y, z, U_lid)
    
    # Apply velocity boundary conditions
    apply_velocity_bc_3d!(u, v, w, domain, x, y, z, U_lid)
    
    # Apply vorticity boundary conditions  
    apply_vorticity_bc_3d!(ω, ψ, u, v, w, domain, x, y, z, U_lid, h)
    
    # Set solid regions
    apply_solid_region_conditions_3d!(ψ, ω, u, v, w, domain)
end

"""
    apply_vector_stream_function_bc_3d!(ψ::Array{Float64,4}, domain::LShapeDomain{3},
                                        x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                                        U_lid::Float64)

Apply boundary conditions for 3D vector stream function.
For the moving lid: ψz = x + 1, ψx = ψy = 0
For other boundaries: ψx = ψy = ψz = 0
"""
function apply_vector_stream_function_bc_3d!(ψ::Array{Float64,4}, domain::LShapeDomain{3},
                                            x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                                            U_lid::Float64)
    nx, ny, nz = size(ψ)[1:3]
    j_top = ny  # Top boundary index
    
    @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
        if domain.boundary_mask[i, j, k]
            if is_moving_lid_node(x[i], y[j], z[k], Val(3))
                # Moving lid: special stream function values
                ψ[i, j, k, 1] = 0.0        # ψx = 0
                ψ[i, j, k, 2] = 0.0        # ψy = 0  
                ψ[i, j, k, 3] = x[i] + 1.0 # ψz = x + 1 (gives u = U_lid at lid)
            else
                # No-slip walls: ψ = 0
                ψ[i, j, k, 1] = 0.0
                ψ[i, j, k, 2] = 0.0
                ψ[i, j, k, 3] = 0.0
            end
        end
    end
end

"""
    apply_velocity_bc_3d!(u, v, w, domain, x, y, z, U_lid)

Apply velocity boundary conditions for 3D case.
"""
function apply_velocity_bc_3d!(u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                              domain::LShapeDomain{3},
                              x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                              U_lid::Float64)
    nx, ny, nz = size(u)
    
    @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j, k]
            # Solid region
            u[i, j, k] = 0.0
            v[i, j, k] = 0.0
            w[i, j, k] = 0.0
        elseif domain.boundary_mask[i, j, k]
            if is_moving_lid_node(x[i], y[j], z[k], Val(3))
                # Moving lid
                u[i, j, k] = U_lid
                v[i, j, k] = 0.0
                w[i, j, k] = 0.0
            else
                # No-slip walls
                u[i, j, k] = 0.0
                v[i, j, k] = 0.0
                w[i, j, k] = 0.0
            end
        end
    end
end

"""
    apply_vorticity_bc_3d!(ω, ψ, u, v, w, domain, x, y, z, U_lid, h)

Apply vorticity boundary conditions for 3D case using velocity gradients.
"""
function apply_vorticity_bc_3d!(ω::Array{Float64,4}, ψ::Array{Float64,4},
                               u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                               domain::LShapeDomain{3},
                               x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                               U_lid::Float64, h::Float64)
    nx, ny, nz = size(u)
    
    @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j, k] || !domain.boundary_mask[i, j, k]
            continue
        end
        
        # Compute boundary vorticity from velocity boundary conditions
        # ω = curl(u) at the boundary, using no-slip conditions
        ωx, ωy, ωz = compute_boundary_vorticity_3d(i, j, k, u, v, w, 
                                                   x, y, z, U_lid, h, domain)
        
        ω[i, j, k, 1] = ωx
        ω[i, j, k, 2] = ωy  
        ω[i, j, k, 3] = ωz
    end
end

"""
    compute_boundary_vorticity_3d(i, j, k, u, v, w, x, y, z, U_lid, h, domain) 
    -> (Float64, Float64, Float64)

Compute vorticity components at a 3D boundary node.
"""
function compute_boundary_vorticity_3d(i::Int, j::Int, k::Int,
                                      u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                      x::Vector{Float64}, y::Vector{Float64}, z::Vector{Float64},
                                      U_lid::Float64, h::Float64, domain::LShapeDomain{3})
    nx, ny, nz = size(u)
    
    # Initialize vorticity components
    ωx = ωy = ωz = 0.0
    
    # For simplicity, use finite difference approximation based on no-slip conditions
    # This is a simplified version - full implementation would need careful treatment
    # of each boundary type and normal direction
    
    if is_moving_lid_node(x[i], y[j], z[k], Val(3))
        # Moving lid: main contribution from lid motion
        ωz = 2.0 * U_lid / h  # From ∂u/∂y at lid
        ωx = -2.0 * v[i, j, k] / h  # From no-slip in z-direction
        ωy = 2.0 * w[i, j, k] / h   # From velocity gradients
    else
        # No-slip walls: vorticity from velocity gradients at wall
        # Use one-sided differences toward interior
        
        # ωx = ∂w/∂y - ∂v/∂z
        if j < ny && domain.fluid_mask[i, j+1, k]
            ωx += 2.0 * w[i, j, k] / h  # ∂w/∂y (one-sided)
        end
        if k < nz && domain.fluid_mask[i, j, k+1]  
            ωx -= 2.0 * v[i, j, k] / h  # -∂v/∂z (one-sided)
        end
        
        # ωy = ∂u/∂z - ∂w/∂x  
        if k < nz && domain.fluid_mask[i, j, k+1]
            ωy += 2.0 * u[i, j, k] / h  # ∂u/∂z
        end
        if i < nx && domain.fluid_mask[i+1, j, k]
            ωy -= 2.0 * w[i, j, k] / h  # -∂w/∂x
        end
        
        # ωz = ∂v/∂x - ∂u/∂y
        if i < nx && domain.fluid_mask[i+1, j, k]
            ωz += 2.0 * v[i, j, k] / h  # ∂v/∂x  
        end
        if j < ny && domain.fluid_mask[i, j+1, k]
            ωz -= 2.0 * u[i, j, k] / h  # -∂u/∂y
        end
    end
    
    return ωx, ωy, ωz
end

"""
    apply_solid_region_conditions_3d!(ψ, ω, u, v, w, domain::LShapeDomain{3})

Set 3D fields in solid regions.
"""
function apply_solid_region_conditions_3d!(ψ::Array{Float64,4}, ω::Array{Float64,4},
                                          u::Array{Float64,3}, v::Array{Float64,3}, w::Array{Float64,3},
                                          domain::LShapeDomain{3})
    nx, ny, nz = size(u)
    
    @inbounds for k in 1:nz, j in 1:ny, i in 1:nx
        if !domain.fluid_mask[i, j, k]
            # Solid region
            ψ[i, j, k, 1] = NaN  # ψx
            ψ[i, j, k, 2] = NaN  # ψy
            ψ[i, j, k, 3] = NaN  # ψz
            ω[i, j, k, 1] = NaN  # ωx  
            ω[i, j, k, 2] = NaN  # ωy
            ω[i, j, k, 3] = NaN  # ωz
            u[i, j, k] = 0.0
            v[i, j, k] = 0.0
            w[i, j, k] = 0.0
        end
    end
end